import angr
import pyvex
from collections import deque

class CallParamAnalyzer:
    def __init__(self, binary_path):
        self.project = angr.Project(binary_path, auto_load_libs=False)
    
    def _resolve_expression(self, expr, tmp_defs, visited=None):
        """
        递归解析表达式中的临时变量，扩展到基本块范围
        返回完全解析后的表达式
        """
        if visited is None:
            visited = set()
        
        # 如果表达式是临时变量，递归解析它
        if isinstance(expr, pyvex.IRExpr.RdTmp):
            tmp = expr.tmp
            if tmp in visited:
                # 避免循环依赖
                return expr
            visited.add(tmp)
            
            if tmp in tmp_defs:
                # 递归解析临时变量的定义
                return self._resolve_expression(tmp_defs[tmp], tmp_defs, visited)
            return expr
        
        # 处理二元操作（如Add64, Sub64）
        elif isinstance(expr, pyvex.IRExpr.Binop):
            resolved_args = [
                self._resolve_expression(arg, tmp_defs, visited)
                for arg in expr.args
            ]
            # 创建新的Binop表达式，使用解析后的参数
            return pyvex.IRExpr.Binop(expr.op, resolved_args)
        
        # 处理一元操作（如32Uto64）
        elif isinstance(expr, pyvex.IRExpr.Unop):
            resolved_arg = self._resolve_expression(expr.args[0], tmp_defs, visited)
            # 创建新的Unop表达式，使用解析后的参数
            return pyvex.IRExpr.Unop(expr.op, resolved_arg)
            
        
        # 处理内存加载操作（如LDle:I32）
        elif isinstance(expr, pyvex.IRExpr.Load):
            resolved_addr = self._resolve_expression(expr.addr, tmp_defs, visited)
            # 创建新的Load表达式，使用解析后的地址
            return pyvex.IRExpr.Load(expr.end, expr.ty, resolved_addr)
        
        # 处理寄存器获取操作
        elif isinstance(expr, pyvex.IRExpr.Get):
            return expr
        
        # 处理常量
        elif isinstance(expr, pyvex.IRExpr.Const):
            return expr
        
        # 其他未处理的表达式类型
        return expr
    
    def _pre_process_statement(self, block, stmt, tyenv, tmp_alias, reg_alias, loop_flag):
        """
        pre process a stmt, and get the stmt's info.
        """
        if isinstance(stmt, pyvex.stmt.Store):

            stmt_data = stmt.data
            st_size = stmt_data.result_size(tyenv)

            # STle(t2) = t3 or STle(0x46ed58) = t3
            if isinstance(stmt_data, pyvex.expr.RdTmp):
                data = 't%d' % (stmt_data.tmp)
                # datas = []
                # datas.append((data, 0))

                if loop_flag:
                    with_op = self._addr_tmp_with_binop(data, tmp_alias)
                    datas = (data, with_op)
                    # alias_results = []
                    # self._find_addr_tmp_alias(data, tmp_alias, alias_results)
                    # print("loop find: %s" % (alias_results))

                    # for sym, offset in alias_results:
                    #     datas.append((sym, offset))

                else:
                    datas = (data, )

                # datas = tuple(datas)

            # STle(t2) = 0x45 or STle(0x46ed58) = 0x7543bd
            elif isinstance(stmt_data, pyvex.expr.Const):
                data = stmt_data.con.value
                datas = (data, )
                # datas = ((data, 0),)

            else:
                l.info("Not support the stmt: %s" % (stmt))
                return None

            if isinstance(stmt.addr, pyvex.expr.RdTmp):
                addr = 't%d' % (stmt.addr.tmp)
                addrs = [('+', (addr, 0))]

                alias_results = []
                self._find_addr_tmp_alias(addr, tmp_alias, alias_results)
                # print("find store addr alias: %s" % (alias_results))

                for t in alias_results:
                    addrs.append(t)

                addrs = tuple(addrs)

            elif isinstance(stmt.addr, pyvex.expr.Const):
                addrs = stmt.addr.con.value

            else:
                l.info("Not support the stmt: %s" % (stmt))
                return None

            action = ('s', addrs, datas, st_size)

        elif isinstance(stmt, pyvex.stmt.WrTmp):

            dst = 't%d' % (stmt.tmp)
            stmt_data = stmt.data
            wr_size = stmt_data.result_size(tyenv)

            # t4 = LDle(t5) or t4 = LDle(0x46ed58)
            if isinstance(stmt_data, pyvex.expr.Load):
                wr_type = 'wl'

                load_addr = stmt_data.addr
                if isinstance(load_addr, pyvex.expr.RdTmp):
                    ld_tmp = 't%d' % (load_addr.tmp)
                    src = [('+', (ld_tmp, 0))]

                    alias_results = []
                    self._find_addr_tmp_alias(ld_tmp, tmp_alias, alias_results)
                    # print("find load addr alias: %s" % (alias_results))

                    for t in alias_results:
                        src.append(t)

                    src = tuple(src)

                elif isinstance(load_addr, pyvex.expr.Const):
                    src = load_addr.con.value

                else:
                    l.info("Not support the stmt: %s" % (stmt))
                    return None

            # t4 = Add(t3, 0x20)
            elif isinstance(stmt_data, pyvex.expr.Binop):
                wr_type = 'wo'
                binop = stmt_data.op

                # TODO
                # if choose_binop and binop not in self.choose_binop:
                #     return None

                opnds = []
                for child_expr in stmt_data.child_expressions:
                    if isinstance(child_expr, pyvex.expr.RdTmp):
                        opnd = 't%d' % (child_expr.tmp)

                    elif isinstance(child_expr, pyvex.expr.Const):
                        opnd = child_expr.con.value

                    else:
                        l.info("Not support the stmt: %s" % (stmt))
                        return None

                    opnds.append(opnd)

                alias_opnds = []
                alias_opnds.append(tuple(opnds))

                if 'Add' in binop or 'Sub' in binop and type(opnds[1]) is int:
                    alias_results = []
                    self._find_tmp_alias(opnds[0], tmp_alias, alias_results)
                    # print("block: %s %s" % (block, stmt))
                    # print("tmp alias: %s" % (tmp_alias))
                    # print("binop alias_results: %s" % (alias_results))

                    # o = opnds[1] if 'Add' in binop else (0-opnds[1])
                    o = opnds[1]
                    for sym in alias_results:
                        alias_opnds.append((sym, o))

                src = (binop, tuple(alias_opnds))

                tmp_alias[dst].append(src)

            # t4 = ITE(t1, t2, t3)
            elif isinstance(stmt_data, pyvex.expr.ITE):

                l.info("Should process the stmt %s in future." % (stmt))
                wr_type = 'wi'

                opnds = []
                for child_expr in stmt_data.child_expressions:
                    if isinstance(child_expr, pyvex.expr.RdTmp):
                        opnd = 't%d' % (child_expr.tmp)

                    elif isinstance(child_expr, pyvex.expr.Const):
                        opnd = child_expr.con.value

                    else:
                        l.info("Not support the stmt: %s" % (stmt))
                        return None

                    opnds.append(opnd)

                src = tuple(opnds)

            # t14 = 1Uto64(t15)
            elif isinstance(stmt_data, pyvex.expr.Unop):
                wr_type = 'wu'
                binop = stmt_data.op

                opnds = []
                for child_expr in stmt_data.child_expressions:
                    if isinstance(child_expr, pyvex.expr.RdTmp):
                        opnd = 't%d' % (child_expr.tmp)

                    elif isinstance(child_expr, pyvex.expr.Const):
                        opnd = child_expr.con.value

                    else:
                        l.info("Not support the stmt: %s" % (stmt))
                        return None

                    opnds.append(opnd)

                src = (binop, tuple(opnds))

                tmp_alias[dst].append(src)

            # t4 = Get(rdi) or t4 = t6
            else:
                wr_type = 'w'

                if isinstance(stmt_data, pyvex.expr.RdTmp):
                    src = 't%d' % (stmt_data.tmp)

                    tmp_alias[dst].append(src)

                elif isinstance(stmt_data, pyvex.expr.Get):
                    src = 'r%d' % (stmt_data.offset)

                    if src == self.sp_name:
                        block.stack_tmps.add(dst)

                    tmp_alias[dst].append(src)

                elif isinstance(stmt_data, pyvex.expr.Const):
                    src = stmt_data.con.value

                else:
                    l.info("Not support the stmt: %s" % (stmt))
                    return None

            action = (wr_type, dst, src, wr_size)

        elif isinstance(stmt, pyvex.stmt.Put):

            if choose_register and stmt.offset in self.ignore_regs:
                return None

            dst = 'r%d' % (stmt.offset)
            stmt_data = stmt.data
            put_size = stmt_data.result_size(tyenv)

            if isinstance(stmt_data, pyvex.expr.RdTmp):
                src = 't%d' % (stmt_data.tmp)
                # srcs = []
                # srcs.append((src, 0))

                if dst == self.sp_name:
                    block.stack_tmps.add(src)

                if dst in reg_alias:
                    tmp_alias[reg_alias[dst]].remove(dst)

                if loop_flag:
                    with_op = self._addr_tmp_with_binop(src, tmp_alias)
                    srcs = (src, with_op)
                    # alias_results = []
                    # self._find_addr_tmp_alias(src, tmp_alias, alias_results)
                    # print("loop find: %s" % (alias_results))

                    # for sym, offset in alias_results:
                    #     srcs.append((sym, offset))

                else:
                    srcs = (src, )

                tmp_alias[src].append(dst)
                reg_alias[dst] = src

                # srcs = tuple(srcs)

            elif isinstance(stmt_data, pyvex.expr.Const):
                src = stmt_data.con.value
                srcs = (src, )
                # srcs = ((src, 0),)

            else:
                l.info("Not support the stmt: %s" % (stmt))
                return None

            action = ('p', dst, srcs, put_size)

        else:
            return None

        return action

    def _expr_to_string(self, expr):
        """将表达式转换为规范化字符串"""
        if isinstance(expr, pyvex.IRExpr.Binop):
            op_name = expr.op.split("_")[-1]
            arg1 = self._expr_to_string(expr.args[0])
            arg2 = self._expr_to_string(expr.args[1])
            return f"{op_name}({arg1}, {arg2})"
        
        elif isinstance(expr, pyvex.IRExpr.Unop):
            op_name = expr.op.split("_")[-1]
            print("op_name", op_name)
            
            arg = self._expr_to_string(expr.args)
            return f"{op_name}({arg})"
        
        elif isinstance(expr, pyvex.IRExpr.Load):
            # 解析端序信息
            end_str = 'le' if expr.end == 'Iend_LE' else 'be'
            # 解析数据类型（去掉前面的'Ity_'）
            data_type = expr.ty.split('_')[-1]
            addr = self._expr_to_string(expr.addr)
            return f"LD{data_type}{end_str}({addr})"
        
        elif isinstance(expr, pyvex.IRExpr.Get):
            reg_name = self.project.arch.register_names[expr.offset]
            return reg_name
        
        elif isinstance(expr, pyvex.IRExpr.Const):
            return hex(expr.con.value)
        
        elif isinstance(expr, pyvex.IRExpr.RdTmp):
            return f"t{expr.tmp}"
        
        else:
            return str(expr)
    
    def analyze_call_params(self, call_addr):
        """分析指定地址的函数调用参数"""
        block = self.project.factory.block(call_addr)
        tmp_defs = {}
        
        # 构建临时变量定义表（基本块内所有定义）
        for stmt in block.vex.statements:
            if isinstance(stmt, pyvex.IRStmt.WrTmp):
                tmp_defs[stmt.tmp] = stmt.data
        
        params = {}
        # 分析参数寄存器赋值
        for stmt in block.vex.statements:
            if isinstance(stmt, pyvex.IRStmt.Put):
                reg = self.project.arch.register_names[stmt.offset]
                # 根据架构确定调用约定参数寄存器
                if self.project.arch.name == 'AMD64':
                    # Windows x64调用约定前4个参数寄存器
                    if reg in ['rcx', 'rdx', 'r8', 'r9']:
                        resolved_expr = self._resolve_expression(stmt.data, tmp_defs)
                        params[reg] = resolved_expr
                elif self.project.arch.name == 'MIPS32' or self.project.arch.name == 'MIPS64':
                    # MIPS调用约定前4个参数寄存器
                    if reg in ['$a0', '$a1', '$a2', '$a3']:
                        resolved_expr = self._resolve_expression(stmt.data, tmp_defs)
                        params[reg] = resolved_expr
                # 可以添加其他架构的处理逻辑
        
        return params
    
    def compare_call_params(self, addr1, addr2):
        """
        比较两个调用点的参数
        返回: (params1, params2, common_elements, has_intersection)
        """
        params1 = self.analyze_call_params(addr1)
        params2 = self.analyze_call_params(addr2)
        
        # 转换为字符串形式比较
        str_params1 = {reg: self._expr_to_string(expr) for reg, expr in params1.items()}
        str_params2 = {reg: self._expr_to_string(expr) for reg, expr in params2.items()}
        
        # 比较值而不仅是引用
        values1 = set(str_params1.values())
        values2 = set(str_params2.values())
        
        common_elements = values1.intersection(values2)
        has_intersection = len(common_elements) > 0
        
        return {
            'addr1_params': str_params1,
            'addr2_params': str_params2,
            'common_elements': list(common_elements),
            'has_intersection': has_intersection
        }



# 使用示例
if __name__ == "__main__":
    analyzer = CallParamAnalyzer("./data/ash")
    
    # 用户输入两个地址
    addr1 = 0x00033a44
    addr2 = 0x00033ac4
    
    # 比较参数
    result = analyzer.compare_call_params(addr1, addr2)
    
    # 输出结果
    print(f"分析地址 0x{addr1:x} 的参数:")
    for reg, expr in result['addr1_params'].items():
        print(f"  {reg}: {expr}")
    
    print(f"\n分析地址 0x{addr2:x} 的参数:")
    for reg, expr in result['addr2_params'].items():
        print(f"  {reg}: {expr}")
    
    if result['has_intersection']:
        print("\n✅ 参数有交集:")
        for elem in result['common_elements']:
            print(f"  {elem}")
    else:
        print("\n❌ 参数无交集")